
% this is the corrected routine that calculates the vertical wavenumber
% spectra of velocity. 

% Author: Ritabrata Thakur, 25 May 2021

function [psd_avg, psdx_array] = welch_estimate_corrected_rt(data_u,data_v,win,NFFT,Fs)
no_psds=1;
 for i=1:length(data_u(1,:)) %time loop
        
        % I am not segmenting the data in depth because of lesser number of
        % points in the vertical. It makes more sense to do it in time.
        % Right now reading at a fixed location which would need change if
        % I want to average across different latitudes

         data_portion_u = squeeze(data_u(:,i));
         data_portion_u = data_portion_u - mean(data_portion_u);%removing the mean 
%         %from individual portion of time series
         data_portion_u = detrend(data_portion_u, 1, 'omitnan');%removing linear trend
         
         data_portion_v = squeeze(data_v(:,i));
         data_portion_v = data_portion_v - mean(data_portion_v);%removing the mean 
%         %from individual portion of time series
         data_portion_v = detrend(data_portion_v, 1, 'omitnan'); %removing linear trend
                
        %% multiplying by a window function because the DFT assumes that the
        % signal is periodic and hence needs to be multiplied by a window
        % function in time domain (or space domain in case of wavenumber spectra). 
        % See https://holometer.fnal.gov/GH_FFT.pdf or holometer.pdf section 8
        
        win_t = win; %transposing for matrix multiplication
        S = sum(win_t.^2)/NFFT;% Calculate the scaling factor. The factor NFFT added on 4Nov21
        xw_u = data_portion_u.*win_t; % Apply windowing here on the data portion
        xw_v = data_portion_v.*win_t;
        
        %% calculate the discrete Fourier transform of the windowed data
        %calculate the FFT of the resultant windowed data with NFFT number of points
        xdft_u = (fft(xw_u, NFFT)); 
        xdft_v = (fft(xw_v, NFFT));
        
        % take the square of the dft
        psdx_u = xdft_u.*conj(xdft_u); % This is same as abs(xdft).^2;
        psdx_v = xdft_v.*conj(xdft_v); % This is same as abs(xdft).^2;
        
        %psdx_u = xdft_u.^2;
        %psdx_v = xdft_v.^2;
        
        % From the Nyquist theorem it follows that the maximal useful
        % frequency is Fs/2. Whether you choose the positive half after calculating the fft or here 
        % after calculating the psd, the result is same
        % Take only first half of the PSD. If the output is complex, a 
        % two-sided spectra will have to be taken
        psdx_u = psdx_u(1:floor(NFFT/2+1)); 
        psdx_v = psdx_v(1:floor(NFFT/2+1));
        
        % For the following normalisation, refer to
        % https://holometer.fnal.gov/GH_FFT.pdf or holometer.pdf section 9
        % (equation 24). The factor 2 is multiplied in the next step
        %normalising by the weight of the window 
        psdx = (1/S) * (psdx_u + psdx_v);  

        %% One-sided spectra for real DFT. 
        psdx(2:end-1) = 2*psdx(2:end-1); %multiplying all frequencies except 0 and N/2 by 2 as 
        %they are overlapped on top of each other. It does not really
        %matter if you multiply the zero and N/2 with 2 or not, we get the
        %same spectra                    
        psdx_array(no_psds,:) = psdx;
        no_psds = no_psds+1;
    
 end
 
     psd_avg = sum(psdx_array,1)/(no_psds-1);
end
     
      
